package edu.zju.gis.dbfg.model.extractor;

import com.google.gson.Gson;
import com.google.gson.JsonArray;
import org.locationtech.jts.geom.Geometry;
//import com.vividsolutions.jts.geom.Geometry;
import edu.zju.gis.dbfg.model.BaseModel;
import edu.zju.gis.dbfg.model.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.util.EsHelper;
import edu.zju.gis.dbfg.model.util.HDFSHelper;
import edu.zju.gis.dbfg.model.util.SparkSessionType;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.Path;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.geotools.geojson.geom.GeometryJSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.io.Serializable;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @author Hu
 * @date 2019/12/31
 * 将 ES 数据导出到个人空间中（HDFS）
 * TODO:
 * (1) 想办法去掉子类对父类构造函数的复写，利用 builder?
 * (2) 利用标签实现 argument 类的注入，而不是通过形参
 **/
public class DataExtractor extends BaseModel<DataExtractorArgs> {

  private final static Logger logger = LoggerFactory.getLogger(DataExtractor.class);

  public DataExtractor(String[] args) {
      this(DataExtractor.class.getName(), args);
  }

  public DataExtractor(String appName, String[] args) {
      this(SparkSessionType.LOCAL, appName, args);
  }

  public DataExtractor(SparkSessionType type, String appName, String[] args) {
      this(type, appName, new SparkConf(), args);
  }

  public DataExtractor(SparkSessionType type, String appName, SparkConf sc, String[] args) {
    super(type, appName, sc, args);
  }

    public static void main(String[] args) throws IOException {
        System.setProperty("es.set.netty.runtime.available.processors", "false");
        DataExtractor job = new DataExtractor(SparkSessionType.SHELL, "DataExtractor", args);
        job.run();
    }

  @Override
  public void run() throws IOException {

      Dataset<Row> df = this.ss.read()
              .format("es")
              .option("inferSchema", "true").load(this.arg.getInput());

      df.registerTempTable(this.getTableNameFromSql(this.arg.getSql()));
      Dataset<Row> resultDf = this.ss.sqlContext().sql(this.arg.getSql());

      // get result fields after sql
      EsHelper esHelper = EsHelper.getInstanceClient();
      StructType schema = resultDf.schema();
      StructField[] fields = schema.fields();
      Map<String, Object> esSchema = esHelper.getEsSchema(this.arg.getInput());
      Map<String, Object> reservedFieldSchema = new HashMap<>();
      List<String> geoShapeField = new ArrayList<>();
      for (StructField field : fields) {
          Map<String, String> fieldMapping = (Map<String, String>) esSchema.get(field.name());
          reservedFieldSchema.put(field.name(), fieldMapping);
          if (fieldMapping.get("type").equals("geo_shape") || fieldMapping.get("type").equals("geo_point")) {
              geoShapeField.add(field.name());
          }
      }
      esHelper.finish();

      Dataset<String> resultStr = (Dataset<String>) resultDf.toJSON().map(new EsJsonToWktMapFunction(geoShapeField.toArray(new String[]{})), Encoders.STRING());

      resultStr.write().text(this.arg.getOutput() + ".zg");

//    resultDf.toJSON().write().format("ZGisFileDataSource")
//      .option("path", this.arg.getOutput())
//      .option("metadata", metadata)
//      .option("shape_name", String.join(",", geoShapeField)).save();

      // write data metadata to output
      Gson gson = new Gson();
      String metadata = gson.toJson(reservedFieldSchema);
      logger.info("metadata: " + metadata);
      HDFSHelper helper = HDFSHelper.getInstance();
      FSDataOutputStream fos = helper.write(new Path(this.arg.getOutput() + ".zg" + File.separator + "metadata"), false);
      fos.write(metadata.getBytes());
      fos.flush();
      fos.close();

      this.close();
  }

    @Override
    protected void close() {
        super.close();
    }

  /**
   * 从一个简单sql语句中获取表名
   * sql形如：select * from [table name] where ...
   * @param sql
   * @return
   */
  private String getTableNameFromSql(String sql) {
    String regex = "\\bfrom\\s*\\S*";
    Pattern p = Pattern.compile(regex);
    Matcher m = p.matcher(sql);
    if (m.find()) {
      return m.group().replace("from","");
    } else {
      throw new ModelFailedException(DataExtractor.class, "getTableNameFromSql()", "unvalid sql, no target table name", sql);
    }
  }

    private class EsJsonToWktMapFunction implements MapFunction<String, String>, Serializable {

        private String[] shapeName;

        public EsJsonToWktMapFunction(String[] shapeName) {
            this.shapeName = shapeName;
        }

        public boolean containsShapeName(String name) {
            for (String n : shapeName) {
                if (n.equals(name)) return true;
            }
            return false;
        }

        private String formatEsGeoJson(Map<String, Object> esGeojson) {
            Gson gson = new Gson();
            Map<String, Object> featureCollection = new HashMap<>();
            JsonArray features = new JsonArray();
            Map<String, Object> feature = new HashMap<>();
            Map<String, Object> geom = new HashMap<>();
            geom.put("coordinates", esGeojson.get("coordinates"));
            String t = (String) esGeojson.get("type");
            String upperT = t.substring(0, 1).toUpperCase() + t.substring(1);
            geom.put("type", upperT);
            feature.put("type", "Feature");
            feature.put("properties", gson.toJson(new HashMap<>()));
            feature.put("geometry", geom);
            features.add(gson.toJsonTree(feature));
            featureCollection.put("type", "FeatureCollection");
            featureCollection.put("features", features);
            return gson.toJson(featureCollection);
        }

        private String geoJsonToWkt(Map<String, Object> esGeoJson) {
            GeometryJSON g = new GeometryJSON();
            String formattedJson = formatEsGeoJson(esGeoJson);
            Geometry geometry = null;
            try {
//                geometry = g.read(formattedJson);
                return geometry.toText();
            } catch (Exception e) {
                logger.error("Unvalid geojson object: " + esGeoJson.toString());
                return null;
            }
        }


        @Override
        public String call(String s) throws Exception {
            Gson gson = new Gson();
            Map<String, Object> json = gson.fromJson(s, Map.class);
            StringBuffer sb = new StringBuffer();
            for (String key : json.keySet()) {
                if (containsShapeName(key)) {
                    Map<String, Object> esGeoJson = gson.fromJson(json.get(key).toString(), Map.class);
                    String wkt = geoJsonToWkt(esGeoJson);
                    sb.append(wkt + "\t");
                } else {
                    sb.append(json.get(key).toString() + "\t");
                }
            }
            return sb.toString().trim();
        }
  }

}
